package com.github.cyf.waf;

import com.github.cyf.common.model.Rs;
import lombok.Data;
import org.apache.catalina.core.ApplicationPart;
import org.apache.commons.io.FilenameUtils;
import org.apache.tomcat.util.buf.HexUtils;
import org.springframework.util.StreamUtils;
import org.springframework.util.StringUtils;
import org.springframework.util.unit.DataSize;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.Part;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.util.ArrayList;
import java.util.Collection;
import java.util.List;
import java.util.regex.Pattern;

/**
 * @author chenyifan
 * @create 2024-11-13 16:21
 */
@Data
public class MultipartRequestChecker {

    private static final Pattern FILE_NAME_INVALID_PATTERN_WIN = Pattern.compile("[\\\\/:*?\"<>|]");

    private static final String POST_METHOD = "POST";

    public static final String MULTIPART = "multipart/";

    /**
     * 允许上传的文件类型
     */
    private List<String> allowFileType = new ArrayList<>();

    /**
     * 允许上传的文件最大大小
     */
    private long maxFileSize = DataSize.ofMegabytes(1L).toBytes();

    public Rs check(HttpServletRequest request) {
        try {
            Collection<Part> parts = request.getParts();
            for (Part part : parts) {
                ApplicationPart applicationPart = (ApplicationPart) part;
                String submittedFileName = applicationPart.getSubmittedFileName();
                if (!StringUtils.hasText(submittedFileName)) {
                    return Rs.error("文件名为空");
                }

                // 检查文件扩展名是否在允许的列表中
                String extension = FilenameUtils.getExtension(submittedFileName).toLowerCase();
                if (!allowFileType.contains(extension)) {
                    return Rs.error("文件类型不支持");
                }

                // 文件名检测
                if (containsInvalid(submittedFileName)) {
                    return Rs.error("文件名存在非法字符！");
                }

                // 文件内容检查，验证文件魔数
                String magicNumber = readHex28Upper(applicationPart.getInputStream());
                FileType fileType = FileType.getBySuffix(extension);
                if (!magicNumber.startsWith(fileType.getMagicNumber())) {
                    return Rs.error("文件内容与文件类型不匹配");
                }

                // 检查上传文件是否超出允许的大小限制
                if (applicationPart.getSize() > maxFileSize) {
                    return Rs.error("上传的文件大小超过限制！");
                }
            }
        } catch (Exception e) {
            return Rs.error(e.getMessage());
        }
        return Rs.ok();
    }

    /**
     * 文件名中是否包含在Windows下不支持的非法字符，包括： \ / : * ? " &lt; &gt; |
     *
     * @param fileName 文件名（必须不包括路径，否则路径符将被替换）
     * @return 是否包含非法字符
     * @since 3.3.1
     */
    private static boolean containsInvalid(String fileName) {
        return !StringUtils.hasText(fileName) || FILE_NAME_INVALID_PATTERN_WIN.matcher(fileName).find();
    }

    private static String readHex28Upper(InputStream is) throws IOException {
        ByteArrayOutputStream os = new ByteArrayOutputStream();
        StreamUtils.copyRange(is, os, 0, 27);
        byte[] bytes = os.toByteArray();
        return HexUtils.toHexString(bytes).toUpperCase();
    }
}